import os
import asyncio
from dotenv import load_dotenv
from camel.agents import ChatAgent
from camel.models import ModelFactory
from camel.societies import Workforce
from camel.toolkits import MCPToolkit
from camel.types import ModelPlatformType
from camel.tasks import Task
from camel.toolkits.base import FunctionTool  # 导入FunctionTool类型

load_dotenv()
Model_Type = os.getenv("MODEL_TYPE")
Model_Api = os.getenv("ZHIPU_API_KEY")
Model_Url = os.getenv("MODEL_URL")

model = ModelFactory.create(
    model_platform=ModelPlatformType.OPENAI,
    model_type=Model_Type,
    api_key=Model_Api,
    url=Model_Url
)


# 递归展平列表的函数
def flatten_tools(tools):
    """递归展平工具列表，确保最终得到的是FunctionTool对象的列表"""
    result = []
    for item in tools:
        if isinstance(item, list):
            # 如果是列表，递归展平
            result.extend(flatten_tools(item))
        elif isinstance(item, (FunctionTool, callable)):
            # 如果是有效的工具类型，直接添加
            result.append(item)
        else:
            # 打印无效类型的警告
            print(f"警告: 忽略无效的工具类型 {type(item)}")
    return result


async def travel_workforce():
    # 构建协调智能体
    coordinator = ChatAgent(
        model=model,
        system_message='你是一个协调任务的智能体，负责拆解任务并协调执行',
        output_language='zh-CN',
    )

    # 构建任务智能体
    task_agent = ChatAgent(
        model=model,
        system_message='你负责任务的规划和分解，将复杂的任务拆解为可执行的子任务',
        output_language='zh-CN',
    )

    # 创建workforce
    workforce = Workforce(
        '旅游规划小组',
        coordinator_agent=coordinator,
        task_agent=task_agent,
    )

    # 连接mcp服务器，为智能体准备工具
    mcp_toolkit = MCPToolkit(config_path="config/mcp_tools.json", timeout=1200)
    await mcp_toolkit.connect()

    # 获取并处理工具
    raw_tools = mcp_toolkit.get_tools()
    # 递归展平工具列表并过滤无效类型
    processed_tools = flatten_tools(raw_tools)

    # 验证处理后的工具类型
    for tool in processed_tools:
        if not isinstance(tool, (FunctionTool, callable)):
            print(f"错误: 工具 {tool} 不是有效的FunctionTool或可调用对象")

    # 创建工作节点
    planmaker_agent = ChatAgent(
        model=model,
        system_message='''
        你是专业的旅游日程规划专家，擅长根据用户需求制定科学合理的整体行程框架...
        # [重要]：你可以使用duckduckgo工具进行网页搜索，以此获取广泛的信息来制定更为合理，实际的行程安排
        ''',
        output_language='zh-CN',
        tools=processed_tools,  # 使用处理后的工具列表
    )

    detailer_agent = ChatAgent(
        model=model,
        system_message='''
        你是景点深度体验规划专家，负责在已确定的整体行程框架下，细化每个景点的具体游玩方案...
        # [重要]：你可以使用duckduckgo工具进行网页搜索，以此获取广泛的信息
        ''',
        output_language='zh-CN',
        tools=processed_tools,
    )

    tourist_agent = ChatAgent(
        model=model,
        system_message='''
        你是拥有丰富旅行经验的评价改进专家，负责从游客体验角度优化行程方案...
        [重要]：你可以使用duckduckgo工具进行网页搜索，以此获取广泛的游客评价
        ''',
        output_language='zh-CN',
        tools=processed_tools,
    )

    # 将工作节点加入工作链
    workforce.add_single_agent_worker(
        '你是专业的旅游日程规划专家',
        worker=planmaker_agent,
    ).add_single_agent_worker(
        '你是景点深度体验规划专家',
        worker=detailer_agent,
    ).add_single_agent_worker(
        '你是拥有丰富旅行经验的评价改进专家',
        worker=tourist_agent,
    )

    # 启动任务
    task = Task(
        content='制定一个2026年1月份，年轻情侣从中国河南省周口市自驾去河南省洛阳市的三日游规划',
        id='0'
    )

    # 用workforce处理任务
    task = await workforce.process_task(task)  # 注意这里添加了await

    print(task.result)

    # 关闭mcp服务器
    await mcp_toolkit.disconnect()


asyncio.run(travel_workforce())
